import sympy
import sympybotics
import pickle
import scipy.io as sio
import numpy as np
rbtdef = sympybotics.RobotDef('barrett_wam', # robot name
                          [('-pi/2', 0, 0, 'q'),
                           ('pi/2', 0, 0, 'q'),
                           ('-pi/2', 0.045, 0.55, 'q'),
                           ('pi/2', -0.045, 0, 'q'),
                           ('-pi/2', 0, 0.3, 'q'),
                           ( 'pi/2', 0, 0, 'q'),
                           ('0',0, 0.06, 'q')], # (alpha, a, d, theta)
                            dh_convention='standard' # either 'standard' or 'modified'
                              )
rbtdef.frictionmodel = {'Coulomb', 'viscous'} # options are None or a combination of 'Coulomb', 'viscous' and 'offset'
rbtdef.dynparms()
rbt = sympybotics.RobotDynCode(rbtdef, verbose=True)
tau_str = sympybotics.robotcodegen.robot_code_to_func('jl', rbt.invdyn_code, 'tau_out', 'tau', rbtdef)
#print(tau_str)
f = open('barrett/barrett_wam_tau_out.cpp','w')
f.write(tau_str)
f.close()


rbt.calc_base_parms()

print(rbt.dyn.baseparms)
regressor_str = sympybotics.robotcodegen.robot_code_to_func('jl', rbt.H_code, 'regressor_out', 'regressor', rbtdef)
#print(regressor_str)
f = open('barrett/barrett_wam_regressor.cpp','w')
f.write(regressor_str)
f.close()
output_file = open('barrett/barrett_regressor_baseparameter.pkl', 'wb')
pickle.dump(rbt, output_file)
output_file1 = open('barrett/barrett_dh.pkl', 'wb')
pickle.dump(rbtdef, output_file1)


sio.savemat('barrett/barrett_Pb.mat', {'Pb': rbt.dyn.Pbnum})
sio.savemat('barrett/barrett_Pd.mat', {'Pd': rbt.dyn.Pdnum})
sio.savemat('barrett/barrett_Kb.mat', {'Kd': rbt.dyn.Kdnum})

